from collections import deque
from typing import Any, Iterable, Iterator, TypeVar

from zkl_aiutils_datasets.basics import Dataset, DatasetIterator, SkipExhausted, pause, resume, skip
from zkl_aiutils_datasets.basics.iterator import pause_or_error
from zkl_aiutils_datasets.wrapping import wrap_dataset
from .sliced import SlicedDataset

AnyItem = TypeVar('AnyItem')
AnyStateOrError = tuple[Any, Any]


class StridingBatchedIterator(DatasetIterator[tuple[AnyItem, ...]]):
    def __init__(self,
        iterator: Iterator[AnyItem],
        batch_items_n: int,
        stride_items_n: int,
        keep_remainder: bool = False,
    ):
        if not stride_items_n <= batch_items_n:
            raise ValueError(f'Expected stride_items_n<=batch_items_n, got {stride_items_n}.')

        self._iterator = iterator
        self._batch_items_n = batch_items_n
        self._stride_items_n = stride_items_n
        self._keep_remainder = keep_remainder

        self._buffer_items: deque[AnyItem] = deque(maxlen=batch_items_n)
        self._buffer_heads_state_or_error: deque[AnyStateOrError] = deque(maxlen=batch_items_n)

    @property
    def iterator(self) -> Iterator[AnyItem]:
        return self._iterator

    @property
    def batch_items_n(self) -> int:
        return self._batch_items_n

    @property
    def stride_items_n(self) -> int:
        return self._stride_items_n

    @property
    def keep_remainder(self) -> bool:
        return self._keep_remainder

    def _fill_buffer(self, till_items_n: int):
        while len(self._buffer_items) < till_items_n:
            head_state_or_error = None
            head = len(self._buffer_items) % self._stride_items_n == 0
            if head:
                head_state_or_error = pause_or_error(self._iterator)

            item = next(self._iterator)
            self._buffer_items.append(item)

            if head:
                self._buffer_heads_state_or_error.append(head_state_or_error)

    def _pop_buffer(self, pop_batches_n: int):
        for _ in range(pop_batches_n * self._stride_items_n):
            self._buffer_items.popleft()
        for _ in range(pop_batches_n):
            self._buffer_heads_state_or_error.popleft()

    def __next__(self) -> tuple[AnyItem, ...]:
        try:
            self._fill_buffer(till_items_n=self._batch_items_n)
        except StopIteration as err:
            if not (self._keep_remainder and len(self._buffer_items) > 0):
                raise StopIteration from err
        batch_items = tuple(self._buffer_items)
        self._pop_buffer(pop_batches_n=1)
        return batch_items

    def __skip__(self, skip_batches_n: int):
        if skip_batches_n == 0:
            return

        # move buffer head
        skip_items_n = skip_batches_n * self._stride_items_n - len(self._buffer_items)
        if skip_items_n < 0:
            self._pop_buffer(pop_batches_n=skip_batches_n)
        else:
            try:
                skip(self._iterator, skip_items_n)
            except SkipExhausted as err:
                skipped_items_n = len(self._buffer_items) + err.skipped_samples_n
                skipped_batches_n = _compute_batches_n(
                    items_n=skipped_items_n,
                    batch_items_n=self._batch_items_n,
                    stride_items_n=self._stride_items_n,
                    keep_remainder=self._keep_remainder)
                raise SkipExhausted(skipped_batches_n) from err
            finally:
                self._buffer_items.clear()
                self._buffer_heads_state_or_error.clear()

        # move buffer tail
        try:
            self._fill_buffer(till_items_n=self._batch_items_n - self._stride_items_n)
        except StopIteration as err:
            skipped_items_n = skip_items_n + len(self._buffer_items)
            skipped_batches_n = _compute_batches_n(
                items_n=skipped_items_n,
                batch_items_n=self._batch_items_n,
                stride_items_n=self._stride_items_n,
                keep_remainder=self._keep_remainder)
            raise SkipExhausted(skipped_batches_n) from err

    def __pause__(self) -> Any:
        if len(self._buffer_heads_state_or_error) == 0:
            return pause(self._iterator)
        state, error = self._buffer_heads_state_or_error[0]
        if error is not None:
            raise error
        return state


class StridingBatchedDataset(Dataset[tuple[AnyItem, ...]]):
    def __init__(self,
        iterable: Iterable[AnyItem],
        batch_items_n: int,
        stride_items_n: int,
        keep_remainder: bool = False,
    ):
        if not stride_items_n <= batch_items_n:
            raise ValueError(f'Expected stride_items_n<=batch_items_n, got {stride_items_n}.')

        self._iterable = iterable
        self._dataset = wrap_dataset(iterable)
        self._batch_items_n = batch_items_n
        self._stride_items_n = stride_items_n
        self._keep_remainder = keep_remainder

    @property
    def iterable(self) -> Iterable[AnyItem]:
        return self._iterable

    @property
    def batch_items_n(self) -> int:
        return self._batch_items_n

    @property
    def stride_items_n(self) -> int:
        return self._stride_items_n

    @property
    def keep_remainder(self) -> bool:
        return self._keep_remainder

    def __len__(self) -> int:
        return _compute_batches_n(
            items_n=len(self._dataset),
            batch_items_n=self._batch_items_n,
            stride_items_n=self._stride_items_n,
            keep_remainder=self._keep_remainder)

    def __getitem__(self, batch_i: int) -> tuple[AnyItem, ...]:
        if batch_i < 0:
            raise IndexError(f'Expected 0<=batch_i, got {batch_i}.')

        try:
            items_n = len(self._dataset)
        except TypeError:
            items_n = None

        if items_n is not None:
            batches_n = _compute_batches_n(
                items_n=items_n,
                batch_items_n=self._batch_items_n,
                stride_items_n=self._stride_items_n,
                keep_remainder=self._keep_remainder)
            if batch_i >= batches_n:
                raise IndexError(f'Expected batch_i<{batches_n}, got {batch_i}.')

        head = batch_i * self._stride_items_n
        tail = head + self._batch_items_n
        tail = tail if items_n is None else min(tail, items_n)
        return tuple(SlicedDataset(self._dataset, slice(head, tail)))

    def __iter__(self) -> StridingBatchedIterator[AnyItem]:
        return StridingBatchedIterator(
            iterator=iter(self._iterable),
            batch_items_n=self._batch_items_n,
            stride_items_n=self._stride_items_n,
            keep_remainder=self._keep_remainder)

    def __resume__(self, state: Any | None = None) -> StridingBatchedIterator[AnyItem]:
        return StridingBatchedIterator(
            iterator=resume(self._iterable, state),
            batch_items_n=self._batch_items_n,
            stride_items_n=self._stride_items_n,
            keep_remainder=self._keep_remainder)


def _compute_batches_n(*,
    items_n: int,
    batch_items_n: int,
    stride_items_n: int,
    keep_remainder: bool
) -> int:
    if items_n < batch_items_n:
        if keep_remainder and items_n > 0:
            return 1
        return 0
    items_n -= batch_items_n
    batches_n = 1 + items_n // stride_items_n
    if keep_remainder and items_n % stride_items_n:
        batches_n += 1
    return batches_n
